# -*- coding: utf-8 -*-
# created on 2017/4/12
# 

from mathsolver.functions.base import *
from mathsolver.functions.base.base import new_latex
from sympy.abc import x, y, n, a, b
from mathsolver.functions.erxiangshi.basic import get_mul_coeff


# 求二项式的系数和
class EXSSumCoeffs(BaseFunction):
    def solver(self, *args):  # in: a,b,n; out: BaseNumber(sum)
        if not args:
            erxiangshi_exs = self.search(str('exs'))
            try:
                e_a = erxiangshi_exs.a
                e_b = erxiangshi_exs.b
                e_n = erxiangshi_exs.n
                sym_a = e_a.free_symbols
                sym_b = e_b.free_symbols
                syms = sym_a.union(sym_b)
                sym_d = {}
                for sym in syms:
                    sym_d[sym] = 1
                    self.steps.append(["", "令%s = 1, 则" % new_latex(sym)])
                sub_a = e_a.subs(sym_d)
                sub_b = e_b.subs(sym_d)
                poly = BaseNumber((sub_a + sub_b) ** e_n)
                self.steps.append(["", "多项式的系数和为：%s" % poly.printing()])
                self.output.append(poly)
            except Exception:
                new_poly = erxiangshi_exs.poly
                self.steps.append(["", "令%s = 1, 则" % new_latex(x)])
                new_poly = new_poly.subs({x: 1, y: 1})
                self.steps.append(["", "多项式的系数和为：%s" % new_latex(new_poly)])
                self.output.append(BaseNumber(new_poly))
        self.label.add("求二项式的系数和")
        return self


# 多项式系数绝对值和
class ExsAbsSumCoeff(BaseFunction):
    def solver(self, *args):  # in: a,b,n; out: BaseNumber(sum)
        self.label.add("求二项式的系数绝对值的和")
        coeffs = args[0].value2pairs()[0]
        e_a, e_b, e_n = coeffs[a], coeffs[b], coeffs[n]
        sym_a = e_a.free_symbols
        sym_b = e_b.free_symbols
        syms = sym_a.union(sym_b)
        sym_d = {}
        for sym in syms:
            sym_d[sym] = 1
            self.steps.append(["", "令%s = 1" % new_latex(sym)])
        sub_a = e_a.subs(sym_d)
        if sub_a < 0:
            sub_a = -sub_a
        sub_b = e_b.subs(sym_d)
        if sub_b < 0:
            sub_b = -sub_b
        poly = BaseNumber((sub_a + sub_b) ** e_n)
        self.steps.append(["", "∴多项式的系数绝对值的和为：%s" % poly.printing()])
        self.output.append(poly)
        return self


# 偶数项系数和
class ExsEvenSumCoeff(BaseFunction):
    def solver(self, *args):  # in: comm; out: sum
        comm = args[0].value
        pow_dict = comm.as_powers_dict()
        nn = -1
        for k, v in pow_dict.items():
            if isinstance(k, binomial):
                nn = k.args[0]
        res = 0
        for r in range(nn + 1):
            if r & 0x1 == 0:  # 为偶数
                res += get_mul_coeff(comm.subs({Symbol("r"): r}))
        res = BaseNumber(res)
        self.steps.append(["", "二项式偶数项的系数和为：%s" % res.printing()])
        self.output.append(res)
        return self


# 奇数数项系数和
class ExsOddSumCoeff(BaseFunction):
    def solver(self, *args):  # in: comm; out: sum
        comm = args[0].value
        pow_dict = comm.as_powers_dict()
        nn = -1
        for k, v in pow_dict.items():
            if isinstance(k, binomial):
                nn = k.args[0]
        res = 0
        for r in range(nn + 1):
            if r & 0x1 == 1:  # 为奇数
                res += get_mul_coeff(comm.subs({Symbol("r"): r}))
        res = BaseNumber(res)
        self.steps.append(["", "二项式奇数项的系数和为：%s" % res.printing()])
        self.output.append(res)
        return self


# 求二项式系数的和
class EXSSumEXCoeff(BaseFunction):
    def solver(self, *args):  # in: a,b,n; out: 2^n
        ex_sum = BaseNumber(2 ** n)
        self.steps.append(["", "二项式的系数和为：%s" % ex_sum.printing()])
        self.label.add("求二项式系数的和")
        self.output.append(ex_sum)
        return self


# 偶数项或奇数项二项系数和，从1开始
class ExsEvenOrOddExSumCoeff(BaseFunction):
    def solver(self, *args):  # in: (a,b,n), str; out: 2^n
        coeffs = args[0].value2pairs()[0]
        e_a, e_b, e_n = coeffs[a], coeffs[b], coeffs[n]
        ex_sum = BaseNumber(2 ** (e_n - 1))
        if args[1] == "even":
            self.steps.append(["", "二项式偶数项的二项式系数和为:%s" % ex_sum.printing()])
        else:
            self.steps.append(["", "二项式奇数项的二项式系数和为:%s" % ex_sum.printing()])
        self.output.append(ex_sum)
        return self
